import os
import cv2
import rawpy
import numpy as np


def align_images(images):
    # 将第一张图作为参考图像
    ref_image = images[2]
    aligned_images = [ref_image]

    for i in range(1, len(images)):
        # 使用ORB特征检测和匹配
        orb = cv2.ORB_create()
        kp1, des1 = orb.detectAndCompute(ref_image, None)
        kp2, des2 = orb.detectAndCompute(images[i], None)

        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        matches = bf.match(des1, des2)
        matches = sorted(matches, key=lambda x: x.distance)

        # 选择前N个最佳匹配点
        N = 20
        good_matches = matches[:N]

        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

        # 估计变换矩阵
        M, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC)

        # 对当前图像进行对齐
        aligned_image = cv2.warpPerspective(images[i], M, (ref_image.shape[1], ref_image.shape[0]))
        aligned_images.append(aligned_image)

    return aligned_images


def blend_images(images):
    # 计算图像的权重
    weights = np.linspace(0.1, 1, len(images))

    # 将图像加权平均叠加
    blended_image = np.zeros_like(images[0], dtype=np.float32)
    for i in range(len(images)):
        blended_image += images[i] * weights[i]
    blended_image /= np.sum(weights)

    return blended_image.astype(np.uint8)


images_path = "G:\\MyFile\\sky\\630\\cs\\IMG_20230630_024241\\"
images_list = []
names = sorted(os.listdir(images_path))
for i in names:
    if os.path.splitext(i)[1].lower() == '.dng':
        dng_image = rawpy.imread(images_path + i)
        # 转换为RGB图像
        rgb_image = dng_image.postprocess()
        # 将RGB转化为GBR
        # bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
        images_list.append(rgb_image)


# 图像对齐
aligned_images = align_images(images_list)

# 图像叠加
final_image = blend_images(aligned_images)

# 显示结果
cv2.namedWindow("Final Image", 0)
cv2.resizeWindow("Final Image", 1200, 900)
cv2.imshow('Final Image', final_image)
cv2.waitKey(0)
cv2.imwrite('Final Image', final_image)
cv2.destroyAllWindows()
